#!/usr/bin/env python
import os,sys
from wwpy.file.File import File

class Spliter(object):
    """split big files into a lot of small parts."""
    
    def __init__(self,file_name,path=None,node_id=None):
        self.file_name=file_name
        self.file=File(file_name)
        self.size=self.file.size
        self.main_name=self.file.main_name
        self.path=self.file.dir_name if path is None else path
        self.part_path=self._part_path()
        self.node_id=node_id
        
    def _part_path(self):
        if not File.exists(self.path):
            self.file.mkdir(self.path)
        return self.path
    
    def _part_name(self,part_num):
        if self.node_id is not None:
            return self.part_path+os.sep+self.main_name+'-part-'+str(self.node_id)+'-'+str(part_num)
        return self.part_path+os.sep+self.main_name+'-part-'+str(part_num)
    
    def _split_by_size(self,part_size=256):
        '''
        split file by size,default part size is 256M.
        '''
        #1024 * 1024=1048576
        parts=[]
        part_size=part_size * 1048576
        part_num=int(self.size/part_size)+1
        f=open(self.file_name,'rb')
        pn=0
        while pn<=part_num:
            data=f.read(part_size)
            if data=='':
                break
            part_name=self._part_name(pn)
            pf=open(part_name,'wb')
            pf.write(data)
            pf.close()
            parts.append(part_name)
            pn+=1
        f.close()
        return parts
        
    def _split_by_num(self,part_num=10):
        '''
        split file into number parts,default part number is 10. 
        '''
        parts=[]
        part_size=int(self.size/part_num)+1
        f=open(self.file_name,'rb')
        pn=0
        while True:
            data=f.read(part_size)
            if data=="":
                break
            part_name=self._part_name(pn)
            pf=open(part_name,'wb')
            pf.write(data)
            pf.close()
            parts.append(part_name)
            pn+=1
        f.close()
        return parts
    
    def _split_by_line(self,part_lines=1000000):
        '''
        split file by line,default part lines is 1000000.
        '''
        parts=[]
        f=open(self.file_name,'r')
        line_num=0
        part_num=0
        fhs=[]
        pf=None
        while True:
            line=f.readline()
            if len(line)==0:
                break
            line_num+=1
            if line_num<=part_lines:
                if pf not in fhs:
                    part_name=self._part_name(part_num)
                    pf=open(part_name,'w')
                    fhs.append(pf)
                    parts.append(part_name)
                pf.write(line)
            else:
                line_num=0
                part_num+=1
                part_name=self._part_name(part_num)
                pf=open(part_name,'w')
                pf.write(line)
                fhs.append(pf)
                parts.append(part_name)
        f.close()
        for fh in fhs:
            fh.close()
        return parts
        
    def _split_by_key(self,part_num=10,hash_key=True,key_cols=[0],insep=',',outsep=',',filter_rows=None):
        '''
        split file by key,only one part has the key,default part number is 10.
        '''
        from wwpy.data.Hash import Hash
        h=Hash()
        parts=[]
        fhs=[]
        for pn in range(part_num):
            part_name=self._part_name(pn)
            fhs.append(open(part_name,'w'))
            parts.append(part_name)
        f=open(self.file_name,'r')
        line_num=0
        pns=range(len(fhs))
        while True:
            line=f.readline()
            if len(line)==0:
                break
            if filter_rows is not None and line_num in filter_rows:
                continue
            col_fields=line.split(insep)
            key=outsep.join([str(col_fields[key_col]) for key_col in key_cols])
            pn=h.hash_mod(key,part_num,hash_key=hash_key)
            if pn in pns:
                fhs[pn].write(line)
            line_num+=1
        f.close()
        for pn in range(part_num):
            fhs[pn].close()
        return parts
    
    def split(self,by='line',**kwargs):
        try:
            by=by.lower()
            if by=='size':
                return self._split_by_size(**kwargs)
            elif by=='num':
                return self._split_by_num(**kwargs)
            elif by=='line':
                return self._split_by_line(**kwargs)
            elif by=='key':
                return self._split_by_key(**kwargs)
            else:
                raise Exception("No such file split method "+by+"!")
        except Exception,e:
            raise Exception("Split file error:%s" %str(e))
            
if __name__ == '__main__':
    s=Spliter('test')
    s.split(by='NOE')
    
